import timm
import torch.nn as nn
from safetensors import torch
from timm import create_model as cm


class Resnet18(nn.Module):
    def __init__(self):
        super(Resnet18, self).__init__()
        features = cm("resnet18", pretrained=True, features_only=True, )
        delattr(features, "layer4")
        for p in features.parameters():
            p.requires_grad = False

        self.features = features

    def forward(self, x):
        return self.features(x)[1:]


class Resnet34(nn.Module):
    def __init__(self):
        super(Resnet34, self).__init__()
        features = cm("resnet34", pretrained=True, features_only=True)
        delattr(features, "layer4")
        for p in features.parameters():
            p.requires_grad = False
        self.features = features

    def forward(self, x):
        return self.features(x)[1:]


class WideResnet50_2(nn.Module):
    def __init__(self):
        super(WideResnet50_2, self).__init__()
        features = cm("wide_resnet50_2", pretrained=True, features_only=True)
        delattr(features, "layer4")
        for p in features.parameters():
            p.requires_grad = False
        self.features = features

    def forward(self, x):
        return self.features(x)[1:]


def create_model(name='resnet18'):
    if name == "resnet18":
        return Resnet18()
    elif name == "resnet34":
        return Resnet34()
    elif name == "wide_resnet50_2":
        return WideResnet50_2()


if __name__ == '__main__':
    import torch
    model = create_model('wide_resnet50_2')
    input = torch.randn(1, 3, 224, 224)
    output = model(input)
    for i in output:
        print(i.shape)